# utils/round_utils.py
import os
import re
import logging
from typing import List, Dict
from LlmModel.DoubaoModel.DoubaoModel import DoubaoModel
import sys
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)

# 获取logger实例
logger = logging.getLogger(__name__)
logger.info("=== 程序开始执行 ===")
logger = logging.getLogger(__name__)


async def count_discount_rounds_by_llm(conversations: List[Dict[str, str]], base_info: str) -> int:
    """
    使用大模型计算 折扣 讨论轮数
    """
    try:
        # 读取提示词文件 - 修正路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        prompt_file = os.path.join(current_dir, "prompts", "discount_rounds_prompt.txt")

        if not os.path.exists(prompt_file):
            logger.error(f"提示词文件不存在: {prompt_file}")
            return 0

        with open(prompt_file, 'r', encoding='utf-8') as f:
            prompt_template = f.read()

        # 构建提示词
        prompt = prompt_template.replace("{{conversations}}", str(conversations)) \
            .replace("{{base_info}}", str(base_info))

        logger.info(f"完整的提示词为: {prompt}")
        # 创建大模型实例
        doubao_model = DoubaoModel()
        result = await doubao_model.generate_text(prompt)
        logger.info(f"大模型返回折扣讨论轮数: {result}")
        # 提取数字
        match = re.search(r'\d+', result.strip())
        if match:
            rounds = int(match.group())
            logger.info(f"大模型返回折扣讨论轮数: {rounds}")
            return min(rounds, 2)
        else:
            logger.warning(f"无法从大模型响应中提取数字: {result}")
            return 0

    except Exception as e:
        logger.error(f"使用大模型计算折扣轮数失败: {str(e)}")
        return 0


async def data_guarantee_rounds_by_llm(conversations: List[Dict[str, str]], base_info: str) -> int:
    """
    使用大模型计算 数据保证 轮数
    """
    try:
        # 读取提示词文件 - 修正路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        prompt_file = os.path.join(current_dir, "prompts", "data_guarantee_rounds_prompt.txt")
        print("------------------------------1----\n")



        if not os.path.exists(prompt_file):
            logger.error(f"提示词文件不存在: {prompt_file}")
            return 0

        with open(prompt_file, 'r', encoding='utf-8') as f:
            prompt_template = f.read()

        # 构建提示词
        prompt = prompt_template.replace("{{conversations}}", str(conversations)) \
            .replace("{{base_info}}", str(base_info))

        # 创建大模型实例
        doubao_model = DoubaoModel()
        result = await doubao_model.generate_text(prompt)

        # 提取数字
        match = re.search(r'\d+', result.strip())
        if match:
            rounds = int(match.group())
            logger.info(f"大模型返回数据保证讨论轮数: {rounds}")
            return min(rounds, 2)
        else:
            logger.warning(f"无法从大模型响应中提取数字: {result}")
            return 0

    except Exception as e:
        logger.error(f"使用大模型计算数据保证轮数失败: {str(e)}")
        return 0



async def time_discount_rounds_by_llm(conversations: List[Dict[str, str]], base_info: str) -> int:
    """
    使用大模型计算 排竟时长 轮数
    """
    try:
        # 读取提示词文件 - 修正路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        prompt_file = os.path.join(current_dir, "prompts", "time_rounds_prompt.txt")
        print("------------------------------1----\n")



        if not os.path.exists(prompt_file):
            logger.error(f"提示词文件不存在: {prompt_file}")
            return 0

        with open(prompt_file, 'r', encoding='utf-8') as f:
            prompt_template = f.read()

        # 构建提示词
        prompt = prompt_template.replace("{{conversations}}", str(conversations)) \
            .replace("{{base_info}}", str(base_info))

        # 创建大模型实例
        doubao_model = DoubaoModel()
        result = await doubao_model.generate_text(prompt)

        # 提取数字
        match = re.search(r'\d+', result.strip())
        if match:
            rounds = int(match.group())
            logger.info(f"大模型排竟时长讨论轮数: {rounds}")
            return min(rounds, 2)
        else:
            logger.warning(f"无法从大模型响应中提取数字: {result}")
            return 0

    except Exception as e:
        logger.error(f"使用大模型计算排竟时长轮数失败: {str(e)}")
        return 0


async def style_rounds_by_llm(conversations: List[Dict[str, str]], base_info: str) -> int:
    """
    使用大模型计算 创作风格 轮数
    """
    try:
        # 读取提示词文件 - 修正路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        prompt_file = os.path.join(current_dir, "prompts", "style_rounds_prompt.txt")
        print("------------------------------1----\n")



        if not os.path.exists(prompt_file):
            logger.error(f"提示词文件不存在: {prompt_file}")
            return 0

        with open(prompt_file, 'r', encoding='utf-8') as f:
            prompt_template = f.read()

        # 构建提示词
        prompt = prompt_template.replace("{{conversations}}", str(conversations)) \
            .replace("{{base_info}}", str(base_info))

        # 创建大模型实例
        doubao_model = DoubaoModel()
        result = await doubao_model.generate_text(prompt)

        # 提取数字
        match = re.search(r'\d+', result.strip())
        if match:
            rounds = int(match.group())
            logger.info(f"大模型 创作风格 讨论轮数: {rounds}")
            return min(rounds, 2)
        else:
            logger.warning(f"无法从大模型响应中提取数字: {result}")
            return 0

    except Exception as e:
        logger.error(f"使用大模型计算 创作风格 轮数失败: {str(e)}")
        return 0

async def schedule_rounds_by_llm(conversations: List[Dict[str, str]], base_info: str) -> int:
    """
    使用大模型计算 档期 轮数
    """
    try:
        # 读取提示词文件 - 修正路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        prompt_file = os.path.join(current_dir, "prompts", "schedule_rounds_prompt.txt")
        print("------------------------------1----\n")



        if not os.path.exists(prompt_file):
            logger.error(f"提示词文件不存在: {prompt_file}")
            return 0

        with open(prompt_file, 'r', encoding='utf-8') as f:
            prompt_template = f.read()

        # 构建提示词
        prompt = prompt_template.replace("{{conversations}}", str(conversations)) \
            .replace("{{base_info}}", str(base_info))

        # 创建大模型实例
        doubao_model = DoubaoModel()
        result = await doubao_model.generate_text(prompt)

        # 提取数字
        match = re.search(r'\d+', result.strip())
        if match:
            rounds = int(match.group())
            logger.info(f"大模型 档期 讨论轮数: {rounds}")
            return min(rounds, 2)
        else:
            logger.warning(f"无法从大模型响应中提取数字: {result}")
            return 0

    except Exception as e:
        logger.error(f"使用大模型计算 档期 轮数失败: {str(e)}")
        return 0

async def tone_rounds_by_llm(conversations: List[Dict[str, str]], base_info: str) -> int:
    """
    使用大模型计算 调性 轮数
    """
    try:
        # 读取提示词文件 - 修正路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        prompt_file = os.path.join(current_dir, "prompts", "tone_rounds_prompt.txt")
        print("------------------------------1----\n")



        if not os.path.exists(prompt_file):
            logger.error(f"提示词文件不存在: {prompt_file}")
            return 0

        with open(prompt_file, 'r', encoding='utf-8') as f:
            prompt_template = f.read()

        # 构建提示词
        prompt = prompt_template.replace("{{conversations}}", str(conversations)) \
            .replace("{{base_info}}", str(base_info))

        # 创建大模型实例
        doubao_model = DoubaoModel()
        result = await doubao_model.generate_text(prompt)

        # 提取数字
        match = re.search(r'\d+', result.strip())
        if match:
            rounds = int(match.group())
            logger.info(f"大模型 调性 讨论轮数: {rounds}")
            return min(rounds, 2)
        else:
            logger.warning(f"无法从大模型响应中提取数字: {result}")
            return 0

    except Exception as e:
        logger.error(f"使用大模型计算 调性 轮数失败: {str(e)}")
        return 0

